import streamlit as st
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import plotly.express as px
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from config.matplotlib_config import apply_matplotlib_cn


apply_matplotlib_cn()

st.set_page_config(page_title="随机森林", page_icon="🌲")

if not st.session_state.get('logged_in', False):
    st.warning("请先在首页登录以访问此内容")
    st.stop()

st.title("随机森林算法")

tab1, tab2, tab3 = st.tabs(["原理讲解", "代码示例", "试一试"])

with tab1:
    st.header("随机森林算法原理")
    st.markdown("""
        随机森林是一种集成学习方法，通过构建多个决策树并综合它们的预测结果来提高模型的准确性和鲁棒性。

        ### 工作原理
        1. 从原始数据集中进行Bootstrap抽样，生成多个子数据集
        2. 对每个子数据集构建决策树
        3. 在构建决策树时，随机选择特征子集进行分裂
        4. 综合所有决策树的预测结果（分类问题采用投票法，回归问题采用平均法）

        ### 优点
        - 能够处理高维数据，无需特征选择
        - 能够评估特征的重要性
        - 对缺失值和异常值不敏感
        - 不容易过拟合

        ### 缺点
        - 模型解释性较差
        - 训练和预测速度较慢
        - 需要调整的参数较多
        """)

with tab2:
    st.header("随机森林代码示例")
    st.code(
        """
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report

# 加载数据
iris = load_iris()
X, y = iris.data, iris.target

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建随机森林模型
model = RandomForestClassifier(n_estimators=100, random_state=42)

# 训练模型
model.fit(X_train, y_train)

# 预测
y_pred = model.predict(X_test)

# 评估模型
accuracy = accuracy_score(y_test, y_pred)
print(f"准确率: {accuracy:.2f}")
print("分类报告:")
print(classification_report(y_test, y_pred))

# 特征重要性
feature_importance = model.feature_importances_
print("特征重要性:")
for i, importance in enumerate(feature_importance):
    print(f"{iris.feature_names[i]}: {importance:.4f}")
        """,
        language="python",
    )

with tab3:
    st.header("随机森林演示")

    iris = load_iris()
    X, y = iris.data, iris.target
    feature_names = iris.feature_names

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42
    )

    col1, col2 = st.columns(2)

    with col1:
        n_estimators = st.slider("树的数量", 1, 200, 100)
        max_depth = st.slider("最大深度", 1, 20, 10)
        max_features = st.selectbox("最大特征数", ["sqrt", "log2", None])

        model = RandomForestClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            max_features=max_features,
            random_state=42,
        )
        model.fit(X_train, y_train)

        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        st.metric("准确率", f"{accuracy:.4f}")

        st.subheader("特征重要性")
        importance_df = pd.DataFrame(
            {"feature": feature_names, "importance": model.feature_importances_}
        ).sort_values("importance", ascending=True)
        fig_importance = px.bar(
            importance_df, x="importance", y="feature", orientation="h", title="特征重要性"
        )
        st.plotly_chart(fig_importance, use_container_width=True)

    with col2:
        feature_x = st.selectbox(
            "X轴特征", range(len(feature_names)), format_func=lambda x: feature_names[x], index=0
        )
        feature_y = st.selectbox(
            "Y轴特征", range(len(feature_names)), format_func=lambda x: feature_names[x], index=1
        )

        X_selected = X[:, [feature_x, feature_y]]

        X_train_sel, X_test_sel, y_train_sel, y_test_sel = train_test_split(
            X_selected, y, test_size=0.3, random_state=42
        )

        model_sel = RandomForestClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            max_features=max_features,
            random_state=42,
        )
        model_sel.fit(X_train_sel, y_train_sel)

        h = 0.02
        x_min, x_max = X_selected[:, 0].min() - 1, X_selected[:, 0].max() + 1
        y_min, y_max = X_selected[:, 1].min() - 1, X_selected[:, 1].max() + 1
        xx, yy = np.meshgrid(
            np.arange(x_min, x_max, h), np.arange(y_min, y_max, h)
        )
        Z = model_sel.predict(np.c_[xx.ravel(), yy.ravel()])
        Z = Z.reshape(xx.shape)

        fig, ax = plt.subplots(figsize=(10, 8))
        plt.contourf(xx, yy, Z, alpha=0.4, cmap=plt.cm.RdYlBu)
        scatter = plt.scatter(
            X_selected[:, 0], X_selected[:, 1], c=y, s=20, edgecolor='k', cmap=plt.cm.RdYlBu
        )
        plt.colorbar(scatter)
        plt.title("随机森林决策边界")
        plt.xlabel(feature_names[feature_x])
        plt.ylabel(feature_names[feature_y])
        st.pyplot(fig) 